import clickhouse_driver


class ClickhouseConnect:
    def __init__(self):
        self.conn = None
        self.cursor = None
        self._connect()

    def _connect(self):
        """建立数据库连接"""
        try:
            self.conn = clickhouse_driver.connect(
                host='192.168.8.10',
                port=9000,
                user='default',
                password="123456",
                database='keyword'
            )
            self.cursor = self.conn.cursor()
        except Exception as e:
            self._close()  # 确保连接失败时清理资源
            raise RuntimeError(f"连接ClickHouse失败: {str(e)}")

    def _close(self):
        """安全关闭连接"""
        try:
            if self.cursor:
                self.cursor.close()
        except Exception as e:
            print(f"关闭游标异常: {str(e)}")
        finally:
            self.cursor = None

        try:
            if self.conn:
                self.conn.close()
        except Exception as e:
            print(f"关闭连接异常: {str(e)}")
        finally:
            self.conn = None

    def __enter__(self):
        """支持上下文管理协议"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出上下文时自动关闭"""
        self._close()

    # 可选：添加重连机制
    def reconnect(self):
        """重新建立连接"""
        self._close()
        self._connect()

    # 示例查询方法
    def execute_query(self, sql, params=None):
        """执行SQL查询"""
        try:
            if not self.cursor:
                self.reconnect()

            self.cursor.execute(sql, params)
            return self.cursor.fetchall()
        except Exception as e:
            self.reconnect()
            return self.execute_query(sql, params)

if __name__ == '__main__':
    clickhouse = ClickhouseConnect()
    ret = clickhouse.execute_query("select count(*) from day_aba")
    print(ret)